# 目标检测
import argparse
import os

from imageai.Detection import ObjectDetection


def _argparse():
    parse = argparse.ArgumentParser()
    parse.add_argument("--model_name", default='yolov3.pt', action='store', required=False,
                       dest="model_name",
                       help="模型，默认yolov3.pt")
    parse.add_argument("--seperate", default=False, action='store', required=False, dest="seperate",
                       help="分割检测结果，默认为不分割")
    # person,  bicycle,  car, motorcycle, airplane, bus, train,  truck,  boat,  traffic light,  fire hydrant, stop_sign,
    # parking meter,   bench,   bird,   cat,   dog,   horse,   sheep,   cow,   elephant,   bear,   zebra,
    # giraffe,   backpack,   umbrella,   handbag,   tie,   suitcase,   frisbee,   skis,   snowboard,
    # sports ball,   kite,   baseball bat,   baseball glove,   skateboard,   surfboard,   tennis racket,
    # bottle,   wine glass,   cup,   fork,   knife,   spoon,   bowl,   banana,   apple,   sandwich,   orange,
    # broccoli,   carrot,   hot dog,   pizza,   donot,   cake,   chair,   couch,   potted plant,   bed,
    # dining table,   toilet,   tv,   laptop,   mouse,   remote,   keyboard,   cell phone,   microwave,   oven,
    # toaster,   sink,   refrigerator,   book,   clock,   vase,   scissors,   teddy bear,   hair dryer,   toothbrush.
    parse.add_argument("--custom", default='', action='store', required=False, dest="custom",
                       help="检测目标，默认为全部")
    parse.add_argument("--file_path", default='ObjectDetection\\1.jpg', action='store', required=False,
                       dest="file_path",
                       help="图片位置")
    return parse.parse_args()


def main():
    parser = _argparse()

    detector = ObjectDetection()
    if parser.model_name == 'mobilenet_v2-b0353104.pth':
        detector.setModelTypeAsRetinaNet()
    elif parser.model_name == 'tiny-yolov3.pt':
        detector.setModelTypeAsTinyYOLOv3()
    else:
        detector.setModelTypeAsYOLOv3()

    models = "models"
    execution_path = os.getcwd()
    model_path = os.path.join(execution_path, models, parser.model_name)

    check = os.path.exists(model_path)
    if not check:
        raise Exception("model is out of range")
    detector.setModelPath(model_path)
    detector.loadModel()

    custom_objects = None
    if parser.custom != "":
        custom_objects = parser.custom.split(",")

    file_real_path = os.path.join(execution_path, parser.file_path)
    file_name = os.path.basename(file_real_path)
    file_new_path = os.path.join(execution_path, "ObjectDetection\\new_" + file_name)

    if not parser.seperate:
        if custom_objects is not None:
            detections = detector.detectObjectsFromImage(custom_objects=custom_objects,
                                                         input_image=file_real_path,
                                                         output_image_path=file_new_path,
                                                         minimum_percentage_probability=30)
        else:
            detections = detector.detectObjectsFromImage(input_image=file_real_path,
                                                         output_image_path=file_new_path,
                                                         minimum_percentage_probability=30)

        for eachObject in detections:
            print(eachObject["name"], " : ", eachObject["percentage_probability"], " : ", eachObject["box_points"])
            print("--------------------------------")
    else:
        if custom_objects is not None:
            detections, objects_path = detector.detectObjectsFromImage(custom_objects=custom_objects,
                                                                       input_image=file_real_path,
                                                                       output_image_path=file_new_path,
                                                                       minimum_percentage_probability=30,
                                                                       extract_detected_objects=True)
        else:
            detections, objects_path = detector.detectObjectsFromImage(input_image=file_real_path,
                                                                       output_image_path=file_new_path,
                                                                       minimum_percentage_probability=30,
                                                                       extract_detected_objects=True)

        for eachObject, eachObjectPath in zip(detections, objects_path):
            print(eachObject["name"], " : ", eachObject["percentage_probability"], " : ", eachObject["box_points"])
            print("Object's image saved in " + eachObjectPath)
            print("--------------------------------")


if __name__ == '__main__':
    main()
